import os
import sys
import argparse
import json
import torch
import torch.nn as nn
from utils.json import json_to_image
from utils.helper import get_subdirs, set_logger, GPUManager, del_useless_folders
import models as cls_models
import pandas as pd
import glob
import time
from pathlib import Path
import numpy as np
import cv2
import PIL
import platform
import yaml
from timm import create_model
from utils.helper import remove_dataparallel
import ssl
ssl._create_default_https_context = ssl._create_unverified_context
import logging

logger = logging.getLogger('clsModel.Train')
def parse_args():
    """
    Set args parameters
    """
    parser = argparse.ArgumentParser(description='Train a segmentation model.')
    parser.add_argument('--name', default=None,
                        help='Name of this experiment: (default: arch+timestamp)')
    parser.add_argument("--height", type=int, default=224, help="size of image height")
    parser.add_argument("--width", type=int, default=224, help="size of image width")
    parser.add_argument("--model", type=str, default='densenet169', choices=['se_resnext50_32x4d', 'densenet161', 'regnetz_e8.ra3_in1k'], help="Define model name")
    parser.add_argument("--weight", type=str, default=r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/code/adc_classification/ckpts/AA_T6_Vtech_TCVPT/0408_202840_densenet169/densenet169_best_acc.pth', help="Define the location of pretrained weights.")
    parser.add_argument("--save", type=str, default='onnx', help="Define the save_dir.")
    parser.add_argument("--res", type=str, default='T6_vtech_TCVPT_cls_20250408.onnx',help="Define the save name.")
    parser.add_argument("--cls", type=int, default=2, help="number of cls.")
    args = parser.parse_args()

    return args


def main():
    args = parse_args()

    # choose free gpus to train
    # gpu_list = ",".join([str(x) for x in GPUManager().auto_choice(gpu_num=args.gpu_num)])
    # os.environ["CUDA_VISIBLE_DEVICES"] = '6'
    # make save dir
    if not os.path.exists(args.save):
        os.makedirs(args.save, exist_ok=True)
    

    if args.model in ['se_resnext50',  'se_resnext101', 'mobilenet_v2', 'densenet161', \
        'densenet121', 'densenet169', 'densenet201']:
        model = getattr(cls_models, args.model)(num_classes=args.cls, pretrained=None)
    else:
        model = create_model(model_name=args.model, pretrained=False, num_classes=args.cls)
    
    params = torch.load(args.weight, map_location='cpu')
    ymap = params['ymap']
    print(ymap)
    print('code_num: ', len(ymap))
    logger.info('Pretrained from %s has been loaded.' %args.weight)
    model.load_state_dict(remove_dataparallel(params["state_dict"]))
    model.eval()
    
    dummy_input = torch.rand(1, 3, args.height, args.width)
    torch.onnx.export(model, dummy_input, os.path.join(args.save, args.res), input_names = ['input1'], output_names = ['output1'], opset_version=11)

    args.json_res = Path(args.res).with_suffix('.json')
    with open(os.path.join(args.save, args.json_res), 'w') as f:
        b = json.dumps(ymap, indent=4)
        f.write(b)

if __name__ == "__main__":
    main()
